import numpy as np
import matplotlib.pyplot as plt

# 设置中文显示
plt.rcParams["font.family"] = ["WenQuanYi Micro Hei"]
plt.rcParams["axes.unicode_minus"] = False  # 正确显示负号

class CubicSpline:
    """三次样条曲线类，用于根据给定控制点构建平滑曲线"""
    
    def __init__(self, x, y, bc_type='natural', bc_values=None):
        """
        初始化三次样条曲线
        
        参数:
            x: 输入点的x坐标，必须严格递增
            y: 输入点的y坐标
            bc_type: 边界条件类型，'natural'或'clamped'
            bc_values: 边界条件值，如果是'clamped'，则为包含两端导数的元组
        """
        # 检查输入
        if len(x) != len(y):
            raise ValueError("x和y的长度必须相同")
        if len(x) < 2:
            raise ValueError("至少需要2个点来构建样条曲线")
        if not np.all(np.diff(x) > 0):
            raise ValueError("x必须严格递增")
            
        self.x = np.array(x)
        self.y = np.array(y)
        self.n = len(x)  # 点的数量
        self.h = np.diff(x)  # 区间长度
        
        # 计算二阶导数M
        self.M = self._compute_second_derivatives(bc_type, bc_values)
        
        # 计算各段多项式的系数
        self.coefficients = self._compute_coefficients()
    
    def _compute_second_derivatives(self, bc_type, bc_values):
        """计算各点的二阶导数M"""
        n = self.n
        h = self.h
        
        # 创建系数矩阵和右侧向量
        A = np.zeros((n, n))
        B = np.zeros(n)
        
        # 填充中间方程 (i=1到n-2)
        for i in range(1, n-1):
            A[i, i-1] = h[i-1]
            A[i, i] = 2 * (h[i-1] + h[i])
            A[i, i+1] = h[i]
            B[i] = 6 * ((self.y[i+1] - self.y[i])/h[i] - (self.y[i] - self.y[i-1])/h[i-1])
        
        # 处理边界条件
        if bc_type == 'natural':
            # 自然边界条件: M[0] = 0, M[-1] = 0
            A[0, 0] = 1
            B[0] = 0
            A[-1, -1] = 1
            B[-1] = 0
        elif bc_type == 'clamped':
            # 固定边界条件: 已知两端的一阶导数
            if bc_values is None or len(bc_values) != 2:
                raise ValueError("clamped边界条件需要提供两个导数的值")
            
            yp0, ypn = bc_values
            
            # 第一个方程
            A[0, 0] = 2 * h[0]
            A[0, 1] = h[0]
            B[0] = 6 * ((self.y[1] - self.y[0])/h[0] - yp0)
            
            # 最后一个方程
            A[-1, -2] = h[-1]
            A[-1, -1] = 2 * h[-1]
            B[-1] = 6 * (ypn - (self.y[-1] - self.y[-2])/h[-1])
        else:
            raise ValueError("边界条件类型必须是'natural'或'clamped'")
        
        # 求解线性方程组
        return np.linalg.solve(A, B)
    
    def _compute_coefficients(self):
        """计算每段三次多项式的系数 a, b, c, d"""
        coefficients = []
        n = self.n
        
        for i in range(n-1):
            h_i = self.h[i]
            a = self.y[i]
            b = (self.y[i+1] - self.y[i])/h_i - h_i*(2*self.M[i] + self.M[i+1])/6
            c = self.M[i]/2
            d = (self.M[i+1] - self.M[i])/(6*h_i)
            coefficients.append((a, b, c, d))
        
        return coefficients
    
    def __call__(self, x):
        """计算给定点x处的样条曲线值"""
        # 如果输入是单个值，转换为数组
        scalar_input = np.isscalar(x)
        x = np.array(x) if not scalar_input else np.array([x])
        
        # 找到每个x所在的区间
        idx = np.searchsorted(self.x, x) - 1
        # 处理边界情况
        idx = np.clip(idx, 0, self.n-2)
        
        # 计算结果
        result = np.zeros_like(x)
        for i in range(len(x)):
            a, b, c, d = self.coefficients[idx[i]]
            t = x[i] - self.x[idx[i]]
            result[i] = a + b*t + c*t**2 + d*t**3
        
        return result[0] if scalar_input else result

# 示例用法
if __name__ == "__main__":
    # 示例1: 使用自然边界条件
    x = [0, 1, 2, 3, 4]
    y = [0, 1, 0, 1, 0]
    
    # 创建三次样条曲线
    spline_natural = CubicSpline(x, y, bc_type='natural')
    
    # 生成密集点用于绘图
    x_dense = np.linspace(min(x), max(x), 1000)
    y_natural = spline_natural(x_dense)
    
    # 示例2: 使用固定边界条件
    # 指定两端的一阶导数
    spline_clamped = CubicSpline(x, y, bc_type='clamped', bc_values=(1.0, -1.0))
    y_clamped = spline_clamped(x_dense)
    
    # 绘图
    plt.figure(figsize=(10, 6))
    plt.plot(x, y, 'o', label='控制点')
    plt.plot(x_dense, y_natural, '-', label='自然边界条件')
    plt.plot(x_dense, y_clamped, '--', label='固定边界条件 (导数=1和-1)')
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('三次样条曲线插值')
    plt.legend()
    plt.grid(True)
    plt.show()
